package com.rose.datax.dataxadmin.manager.impl;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.rose.datax.dataxadmin.entity.DbConnection;
import com.rose.datax.dataxadmin.entity.MigrationTable;
import com.rose.datax.dataxadmin.entity.MigrationTask;
import com.rose.datax.dataxadmin.entity.Transformer;
import com.rose.datax.dataxadmin.manager.JobFileGenerator;
import com.rose.datax.dataxadmin.mapper.DbConnectionMapper;
import com.rose.datax.dataxadmin.mapper.MigrationTableMapper;
import com.rose.datax.dataxadmin.mapper.MigrationTaskMapper;
import com.rose.datax.dataxadmin.mapper.TransformerMapper;
import com.rose.datax.dataxadmin.utils.ArrayUtils;
import com.zaxxer.hikari.HikariDataSource;
import freemarker.template.Configuration;
import freemarker.template.Template;
import lombok.Cleanup;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.File;
import java.io.FileOutputStream;
import java.io.OutputStreamWriter;
import java.sql.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

@Service
@Slf4j
public class JobFileGeneratorImpl implements JobFileGenerator {

    @Value("${datax.home}")
    private String dataxHome;

    @Resource
    private MigrationTaskMapper taskMapper;
    @Resource
    private MigrationTableMapper tableMapper;

    @Resource
    private DbConnectionMapper connectionMapper;

    @Resource
    private TransformerMapper transformerMapper;

    private HikariDataSource sourceDatasource = null;

    private ThreadLocal<Connection> sourceContainer = new ThreadLocal<>();

    private Configuration freemarkerCfg = null;

    private Configuration getFreemarkerConfiguration(){
        Configuration configuration = new Configuration(Configuration.VERSION_2_3_20);
        configuration.setDefaultEncoding("UTF-8");
        configuration.setClassForTemplateLoading(JobFileGeneratorImpl.class, "/");
        return configuration;
    }

    private void init(DbConnection source){
        sourceDatasource = new HikariDataSource();
        sourceDatasource.setJdbcUrl(source.getJdbcURL());
        sourceDatasource.setUsername(source.getUsername());
        sourceDatasource.setPassword(source.getPassword());
        sourceDatasource.setMaximumPoolSize(10);
        sourceDatasource.setMinimumIdle(5);
    }

    @Override
    public boolean execute(long taskId) {
        boolean result = false;
        try {
            MigrationTask task = taskMapper.selectById(taskId);
            DbConnection sourceConn = connectionMapper.selectById(task.getSourceId());
            sourceConn.setDb(task.getSourceDb());
            DbConnection targetConn = connectionMapper.selectById(task.getTargetId());
            targetConn.setDb(task.getTargetDb());
            init(sourceConn);
            if(null == freemarkerCfg){
                freemarkerCfg = this.getFreemarkerConfiguration();
            }
            long start = System.currentTimeMillis();
            //查询出需要同步的数据库表
            List<MigrationTable> migrationTables = tableMapper.selectList(new QueryWrapper<MigrationTable>().eq("task_id", taskId));
            List<List<MigrationTable>> splitTables = ArrayUtils.splitArray(migrationTables);

            CountDownLatch countDownLatch = new CountDownLatch(splitTables.size());
            //CPU 核数
            int cpuNum = Runtime.getRuntime().availableProcessors();
            ExecutorService executor = new ThreadPoolExecutor(cpuNum, cpuNum * 2, 60000L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<Runnable>());

            for(List<MigrationTable> list: splitTables){
                executor.submit(new JobFileGeneratorThread(taskId, list, sourceConn, targetConn, countDownLatch));
            }
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            executor.shutdown();
            long time = System.currentTimeMillis() - start;
            log.info("job file generator finish! cost time: {}s", time/1000);
            result = true;
        } catch (Exception e) {
            log.error(e.getMessage());
        }finally {
            sourceDatasource.close();
        }
        return result;
    }

    private void generator(long taskId, String tableName, Map<String, Object> data){
        try {
            if(!this.dataxHome.endsWith("/")){
                this.dataxHome = this.dataxHome +"/";
            }
            Template template = freemarkerCfg.getTemplate("job-template.json");
            File folder = new File(String.format("%sjob/task-%s/",dataxHome,taskId+""));
            if(!folder.exists()){
                folder.mkdirs();
            }
            String filePath = String.format("%s/%s.json",folder.getAbsolutePath(),tableName);
            @Cleanup
            FileOutputStream fileOutputStream = new FileOutputStream(new File(filePath));
            @Cleanup
            OutputStreamWriter outputStreamWriter = new OutputStreamWriter(fileOutputStream, "UTF-8");
            template.process(data, outputStreamWriter);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private CharSequence getColumnsString(List<String> columns) {
        StringBuilder stb = new StringBuilder();
        for (String s : columns) {
            stb.append("\"");
            stb.append(s);
            stb.append("\",");
        }
        return stb.subSequence(0, stb.length() - 1);
    }

    /**
     * 根据表的记录数设置对应的datax 通道数
     * @param records
     * @return
     */
    private int getChannelNumber(long records){
        int channel = 4;
        if(records > 10000){
            channel = 8;
        }
        if(records > 100000){
            channel = 10;
        }
        if(records > 1000000){
            channel = 12;
        }
        return channel;
    }

    private Connection getSourceConnection(){
        Connection conn = null;
        try{
            conn = sourceContainer.get();
            if(null == conn || conn.isClosed()) {
                conn = sourceDatasource.getConnection();
                sourceContainer.set(conn);
            }
        }catch (SQLException e){
            e.printStackTrace();;
        }
        return conn;
    }

    private long getSourceTableRows(String tableName){
        String sql = "select count(0) from `"+tableName+"`";
        long count = 0;
        try{
            @Cleanup
            Connection connection = getSourceConnection();
            @Cleanup
            PreparedStatement pstmt = connection.prepareStatement(sql);
            @Cleanup
            ResultSet rs = pstmt.executeQuery();
            if(rs.next()){
                count = rs.getLong(1);
            }
        }catch (SQLException e){
            e.printStackTrace();
        }finally {
            sourceContainer.remove();
        }
        return count;
    }


    /**
     * 表结构SQL
     * @param tableName
     * @return
     */
    private String getCreateTableSQL(String tableName){
        String sql = null;
        try{
            @Cleanup
            Connection connection = getSourceConnection();
            @Cleanup
            Statement stmt = connection.createStatement();
            @Cleanup
            ResultSet rs = stmt.executeQuery("show create table `" + tableName +"`");
            if(rs.next()){
                sql = rs.getString(2);
            }
        }catch (SQLException e){
            e.printStackTrace();
        }finally {
            sourceContainer.remove();
        }
        return sql;
    }

    /**
     * 获取源数据库表的所有列名
     * @param database
     * @param tableName
     * @return
     */
    private List<String> getSourceTableColumns(String database, String tableName){
        List<String> columns = new ArrayList<>();
        String sql = "select column_name from information_schema.columns t where table_schema=? and table_name=?";
        try {
            @Cleanup
            Connection connection = getSourceConnection();
            @Cleanup
            PreparedStatement pstmt = connection.prepareStatement(sql);
            pstmt.setString(1, database);
            pstmt.setString(2, tableName);
            @Cleanup
            ResultSet rs = pstmt.executeQuery();
            if(null != rs){
                while (rs.next()){
                    columns.add(rs.getString(1));
                }
            }
        }catch (SQLException e){
            e.printStackTrace();
        }finally {
            sourceContainer.remove();
        }
        return columns;
    }

    @Override
    public List<String> getColumns(long taskId, String tableName){
        MigrationTask task = taskMapper.selectById(taskId);
        DbConnection sourceConn = connectionMapper.selectById(task.getSourceId());
        sourceConn.setDb(task.getSourceDb());
        init(sourceConn);
        return getSourceTableColumns(task.getSourceDb(), tableName);
    }

    /**
     * 查询表的主键
     * @param tableName
     * @return
     */
    private List<String> getTablePrimaryKeys(String database, String tableName){
        List<String> primaryKeys = new ArrayList<>();
        String sql = "select column_name from information_schema.columns t where table_schema=? and column_key='pri' and table_name=?";
        try{
            @Cleanup
            Connection connection = getSourceConnection();
            @Cleanup
            PreparedStatement pstmt = connection.prepareStatement(sql);
            pstmt.setString(1, database);
            pstmt.setString(2, tableName);
            @Cleanup
            ResultSet rs = pstmt.executeQuery();
            if(null != rs){
                while (rs.next()){
                    primaryKeys.add(rs.getString(1));
                }
            }
            log.info("get table primary keys, {}", primaryKeys);
        }catch (SQLException e){
            e.printStackTrace();
        }finally {
            sourceContainer.remove();
        }
        return primaryKeys;
    }

    class JobFileGeneratorThread extends Thread{

        private Long taskId;
        private List<MigrationTable> tables;
        private CountDownLatch countDownLatch;
        private DbConnection source;
        private DbConnection target;

        private JobFileGeneratorThread(Long taskId,List<MigrationTable> tables,
                                       DbConnection source, DbConnection target,CountDownLatch countDownLatch){
            this.taskId = taskId;
            this.tables = tables;
            this.source = source;
            this.target = target;
            this.countDownLatch = countDownLatch;
        }

        @Override
        public void run(){
            for(MigrationTable table: tables){
                String tableName = table.getTableName();
                table.setRows(getSourceTableRows(tableName));

                String sql = getCreateTableSQL(tableName);
                // 更新表
                tableMapper.updateById(table);
                //表的列名
                List<String> columns = getSourceTableColumns(source.getDb(), tableName);
                List<String> primaryKeys = getTablePrimaryKeys(source.getDb(), tableName);
                String pk = "";
                if (primaryKeys.size() > 0) {
                    pk = primaryKeys.get(0);
                }
                List<Transformer> transformers = transformerMapper.selectList(new QueryWrapper<Transformer>().eq("task_id", taskId).eq("table_name",tableName));
                int channel = getChannelNumber(table.getRows());
                Map<String, Object> data = new HashMap<>();
                data.put("table", tableName);
                data.put("channel", channel);
                data.put("columns", getColumnsString(columns));
                data.put("source", source);
                data.put("pk", pk);
                data.put("tableStruct", sql.replaceAll("\n", " "));
                data.put("target", target);
                if(null != transformers && transformers.size() > 0){
                    data.put("transformer", transformers);
                }
                generator(taskId, tableName, data);
            }
            countDownLatch.countDown();
        }

    }

}
